# +----------------------------------------------------------------------
# | ChatWork
# +----------------------------------------------------------------------
# | 这不是一个自由软件,您只能在不用于商业目的的前提下对程序代码进行修改和使用。
# | 任何企业和个人不允许对程序代码以任何形式任何目的再发布,商业使用请获取授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import os
import json
import time
from fastapi import UploadFile, Request
from tortoise.expressions import Q

from common.utils.config import ConfigUtil
from config import get_settings
from exception import AppException
from common.utils.urls import UrlUtil
from common.utils.times import TimeUtil
from common.utils.tools import ToolsUtil
from common.utils.files import FilesUtil
from common.enums.voice import VoiceEnum
from common.models.users import UserModel
from common.models.agents import AgentRobotModel, AgentRobotInstructModel
from common.models.agents import AgentAttachmentModel
from common.models.agents import AgentChatRecordsModel
from common.models.agents import AgentRobotPublishModel
from plugins.storage.driver import StorageDriver
from apps.api.schemas.agent.chat import dialog_schema as schema
from apps.api.schemas.index_schema import UploadResultVo


class ChatDialogService:

    @classmethod
    async def history(cls, params: schema.ChatHistorySearchIn, query_params: dict):
        """
        智能体对话记录

        Args:
            params (schema.ChatHistorySearchIn): 搜索参数
            query_params (dict): 检索参数

        Returns:
           List[any]: 对话记录列表

        Author:
           zero
        """
        # 发布的参数
        user_id: int = query_params["user_id"]
        publish_id: int = query_params["publish_id"]
        publish_uid: str = query_params["publish_uid"]
        publish_scene: str = query_params["publish_scene"]

        # 查询的条件
        where: list = []
        if publish_id:
            where.append(Q(scene=publish_scene))
            where.append(Q(publish_id=publish_id))
            where.append(Q(publish_uid=publish_uid))
            where.append(Q(session_id=0))
        else:
            where.append(Q(scene=params.scene))
            where.append(Q(user_id=user_id))
            where.append(Q(publish_id=0))
            where.append(Q(session_id=params.sid if params.scene == "normal" else 0))

        # 验证智能体
        robot = await AgentRobotModel.filter(code=params.bot).first()
        if not robot:
            raise AppException("智能体不存在!")

        # 查询对话记录
        records = await (AgentChatRecordsModel
                         .filter(*where)
                         .filter(robot_id=robot.id)
                         .filter(is_delete=0)
                         .filter(user_delete=0)
                         .limit(30)
                         .order_by("-id")
                         .all())

        # 查询用户信息
        user_avatar: str = ""
        if user_id:
            user = await UserModel.filter(id=user_id).first().values("avatar") or {}
            user_avatar = await UrlUtil.to_absolute_url(user.get("avatar", ""))

        # 角色转换字典
        role_dict = {"system": "System", "user": "Human", "assistant": "AI"}

        # 循环处理数据
        lists = []
        records = list(reversed(records))
        for item in records:
            quotes = json.loads(item.quotes or "[]") if robot.is_show_quote else []
            context = json.loads(item.context or "[]") if robot.is_show_context else []

            # 处理上下文
            ctx_array = []
            for ctx in context:
                content = ctx["content"]
                if isinstance(ctx["content"], list):
                    annex, value = ("", "")
                    for ctx_item in ctx["content"]:
                        name = ctx_item.get('file_name', '-')
                        if ctx_item["type"] == "image":
                            annex += f"Input an image: {name}\n"
                        elif ctx_item["type"] == "video":
                            annex += f"Input an video: {name}\n"
                        elif ctx_item["type"] == "audio":
                            annex += f"Input an audio: {name}\n"
                        elif ctx_item["type"] == "file":
                            annex += f"Input a file: {name}\n"
                        elif ctx_item["type"] == "text":
                            value += ctx_item["text"] + "\n"

                    annex += "\n" if annex else ""
                    content = (annex + value).strip("\n")
                role = role_dict.get(ctx["role"], ctx["role"])
                ctx_array.append({"role": role, "content": content})

            # 答复附件
            annex = []
            annex_lists = json.loads(item.annex or "[]")
            for att in annex_lists:
                att["file_path"] = await UrlUtil.to_absolute_url(att["file_path"])
                annex.append(att)

            # 提问附件
            attach = []
            attach_lists = json.loads(item.attach or "[]")
            for att in attach_lists:
                att["file_path"] = await UrlUtil.to_absolute_url(att["file_path"])
                attach.append(att)

            # 消耗信息
            chat_model = ""
            usage_lists = json.loads(item.flows_usage or "[]")
            for us in usage_lists:
                if us.get("scene") == "chat":
                    chat_model = us.get("alias")
                    break

            # 猜你想问
            issues = json.loads(item.issues or "[]")

            # MCP工具
            tools_list = json.loads(item.tools_list or "[]")
            for mcp in tools_list:
                mcp["tool_logo"] = await UrlUtil.to_absolute_url(mcp["tool_logo"])

            # 问题
            lists.append(schema.ChatHistoryListVo.Question(
                id=item.id,
                type="question",
                avatar=user_avatar,
                content=item.question,
                create_time=TimeUtil.timestamp_to_date(item.create_time),
                attach=attach
            ))

            # 答复
            lists.append(schema.ChatHistoryListVo.Answer(
                id=item.id,
                type="answer",
                avatar=await UrlUtil.to_absolute_url(robot.logo),
                thinks=item.thinks or "",
                content=item.answer or "",
                quotes=quotes,
                context=ctx_array,
                annex=annex,
                issues=issues,
                tools_list=tools_list,
                web_search=json.loads(item.web_search or "[]"),
                chat_model=chat_model,
                end_clear=bool(item.end_clear),
                is_clear=bool(item.is_clear),
                task_time=item.task_time + "s",
                create_time=TimeUtil.timestamp_to_date(item.create_time)
            ))

        return lists

    @classmethod
    async def clear(cls, post: schema.ChatDialogClearIn, query_params: dict):
        """
        对话记录清空

        Args:
            post (schema.ChatDialogClearIn): 清空参数
            query_params (dict): 检索参数

        Author:
           zero
        """
        user_id: int = query_params['user_id']
        publish_id: int = query_params['publish_id']
        publish_uid: str = query_params['publish_uid']
        publish_scene: str = query_params['publish_scene']

        # 查询的条件
        where: list = []
        if publish_id:
            where.append(Q(session_id=0))
            where.append(Q(scene=publish_scene))
            where.append(Q(publish_id=publish_id))
            where.append(Q(publish_uid=publish_uid))
        else:
            where.append(Q(scene=post.scene))
            where.append(Q(user_id=user_id))
            where.append(Q(publish_id=0))
            where.append(Q(session_id=post.sid if post.scene == "normal" else 0))

        # 验证智能体
        robot = await AgentRobotModel.filter(code=post.bot).first()
        if not robot:
            raise AppException("智能体不存在!")

        # 清空上下文
        await (AgentChatRecordsModel
               .filter(*where)
               .filter(robot_id=robot.id)
               .filter(session_id=post.sid)
               .filter(is_delete=0)
               .filter(user_delete=0)
               .filter(is_clear=0)
               .update(
                    is_clear=1,
                    update_time=int(time.time())
               ))

        # 标记清空节点
        await (AgentChatRecordsModel
               .filter(*where)
               .filter(robot_id=robot.id)
               .filter(session_id=post.sid)
               .filter(is_delete=0)
               .filter(user_delete=0)
               .filter(is_clear=1)
               .order_by("-id")
               .limit(1)
               .update(end_clear=1))

    @classmethod
    async def verify(cls, bot: str, name: str) -> schema.ChatPublishVerifyVo:
        """
        发布密钥校验

        Args:
            bot (str): 智能体编号
            name (str): 发布窗口名称

        Returns:
            schema.ChatPublishVerifyVo: 智能体信息

        Author:
            zero
        """
        # 查智能体
        robot = await (AgentRobotModel
                       .filter(code=bot)
                       .filter(is_delete=0)
                       .first())

        # 语音播报
        broad = json.loads(robot.broadcast or "{}")
        voice = broad.get("model", "").split(":")[0]
        speech = VoiceEnum.get_default_speech(voice) or {}
        broadcast = {
            "model": str(broad.get("model", "")),
            "speed": broad.get("speed", 0.25),
            "speed_range": speech.get("speed_range", [0.25, 4.0, 0.1])
        }

        # 内置指令
        instruct = []
        directs = await AgentRobotInstructModel.filter(robot_id=robot.id, scene="direct").order_by("-sort").all()
        for direct in directs:
            instruct.append({"name": direct.keyword, "content": direct.content})

        # 联网搜索
        is_web_search: bool = bool(await ConfigUtil.get("web_search", "status", 0))

        # 返回结果
        return schema.ChatPublishVerifyVo(
            code=robot.code,
            logo=await UrlUtil.to_absolute_url(robot.logo),
            name=name,
            intro=robot.intro,
            welcome_text=robot.welcome_text or "",
            is_web_search=is_web_search,
            max_upload_num=robot.max_upload_num,
            is_upload_files=robot.is_upload_files,
            is_upload_image=robot.is_upload_image,
            is_show_context=robot.is_show_context,
            is_show_quote=robot.is_show_quote,
            is_show_timed=robot.is_show_timed,
            broadcast=broadcast,
            instruct=instruct
        )

    @classmethod
    async def upload(cls, file_in: UploadFile, user_id: int, ip: str) -> UploadResultVo:
        """
        对话文件上传

        Args:
            file_in (UploadFile): 清空参数
            user_id (int): 用户ID
            ip (str): 来源IP

        Returns:
           UploadResultVo: 上传结果

        Author:
           zero
        """
        supported_extensions = [
            "doc", "docx", "xls", "xlsx", "ppt", "pptx", "pdf", "csv", "txt",
            "jpg", "jpeg", "png",
            "mp4"
        ]

        extensions = file_in.filename.split(".")[-1].lower()
        if extensions not in supported_extensions:
            raise AppException("不被支持的文件扩展")

        file_type = "file"
        if extensions in ["jpg", "jpeg", "png"]:
            file_type = "image"
        elif extensions in ["mp4"]:
            file_type = "video"

        result = await StorageDriver.upload(
            file_in=file_in,
            root=get_settings().UPLOAD.get("root") or "",
            stores=get_settings().UPLOAD.get("path") or "",
            folder="agents/chat",
            synchro=True if file_type == "file" else False
        )

        content = ""
        text_path = ""
        if file_type == "file":
            # 读取内容
            file_path = UrlUtil.to_root_path(result.get("path"))
            content = await FilesUtil.parse_raw(file_path)
            # 保存内容

            save_name = ToolsUtil.make_md5_str(file_path)
            save_date = TimeUtil.timestamp_to_date(int(time.time()), formats="%Y%m%d")
            save_path = UrlUtil.to_runtime_path(f"parsing/chat/{save_date}/{save_name}.md")
            text_path = f"parsing/chat/{save_date}/{save_name}.md"
            await ToolsUtil.write_file(save_path, content)
            # 删除文件
            engine: str = await UrlUtil.get_storage_engine()
            if engine != "local" and os.path.exists(file_path):
                os.remove(file_path)

        md5 = ToolsUtil.make_md5_str(result.get("name").strip() + content.strip())
        fid = ToolsUtil.make_md5_str(result.get("path"))
        await AgentAttachmentModel.create(
            id=fid,
            user_id=user_id,
            file_md5=md5,
            file_type=file_type,
            file_name=result.get("name"),
            file_size=result.get("size"),
            file_ext=result.get("ext"),
            file_path=result.get("path"),
            text_path=text_path,
            ip=ip,
            create_time=int(time.time())
        )

        return UploadResultVo(
            id=fid,
            name=result["name"],
            size=result["size"],
            ext=result["ext"],
            path=result["path"],
            url=await UrlUtil.to_absolute_url(result["path"])
        )

    @classmethod
    async def check_publish_api(cls, request: Request):
        """
        验证发布API

        Args:
            request (Request): 请求对象

        Returns:
            publish|None

        Author:
            zero
        """
        permits: str = request.headers.get("Permits", "").strip()
        secrets: str = request.headers.get("Secrets", "").strip()
        identity: str = request.headers.get("Identity", "").strip()
        if not permits:
            authorization = request.headers.get("Authorization", "")
            authorization = authorization.strip("Bearer ").strip("bearer ")
            if authorization.strip().startswith("sk-"):
                permits = authorization.strip()
        if not identity:
            array = permits.split("@")
            if len(array) >= 2:
                permits = array[0].strip()
                identity = array[1].strip()

        if permits:
            publish = await AgentRobotPublishModel.filter(api_key=permits, is_delete=0).first()
            if not publish:
                raise AppException(msg="apiKey无效", code=1100)
            if publish.api_secret and publish.api_secret != secrets:
                raise AppException(msg="访问密码错误", code=1200)
            if (not identity and publish.scene == "web") or len(identity) > 32:
                raise AppException(msg="身份鉴权失败", code=1300)
            if int(time.time()) >= publish.expire_time > 0:
                raise AppException("apiKey已过期", code=1101)
            if publish.rpm_limit:
                minute = TimeUtil.current_minute_range()
                rpm = await (AgentChatRecordsModel
                             .filter(create_time__gte=minute[0])
                             .filter(create_time__lte=minute[1])
                             .filter(ip=request.client.host)
                             .count())
                if rpm >= publish.rpm_limit:
                    raise AppException(f"已超出每分钟请求次数/{rpm}(RPM),请稍后再试!", code=1102)

            robot = await AgentRobotModel.filter(id=publish.robot_id).first().values("code") or {}
            if not identity:
                identity = "0"

            return {
                "id": publish.id,
                "bot": robot.get("code", ""),
                "name": publish.name,
                "scene": publish.scene,
                "permits": permits,
                "secrets": secrets,
                "identity": identity
            }
        return None
